#include <math.h>
#include <iostream>
#include <fstream>
#include <stdlib.h>
#include <stdio.h>
//#include <mpi.h>
#define A1 -1
#define A2 2
#define B1 -2
#define B2 2
//#define M 40
//#define N 40
#define EPS 0.000001

using namespace std;

double func_u(double x, double y){
    return exp(1-(x+y)*(x+y));
}

double func_q(double x, double y){
	if (x+y < 0)
		return 0;
	else
	    return (x+y)*(x+y);
}

double func_k(double x, double y){
    return 4+x;
}

double func_F(double x, double y){
    return 2*(8 + 3*x + y - 4*(4+x)*(x+y)*(x+y))*func_u(x, y) + func_q(x, y)*func_u(x, y);
}
 
double* vector_diff(double *w1, double *w0, int M, int N){       // Разность двух векторов
    double *res = new double [(M+1)*(N+1)];

    for(int i = 0; i < M+1; i++)
        for (int j = 0; j < N+1; j++)
            if(i == 0 or i == M or j == 0 or j == N)
                res[i*(N+1) + j] = 0;                            // Значение w0 и w1 на границах сопадают
            else
                res[i*(N+1) + j] = w1[i*(N+1) + j] - w0[i*(N+1) + j];

    return res;
}

    double ro_i(int i,int M, int N){                           // Вспомогательные функции для скалярного произведения в пространстве H
        if (i == 1 or i == M-1)
            return 0.25;                                       // Попробуем поставить везде параметр 0.25 вместо 0.5, это существенно уменьшает число итераций
        else
            return 1;
    }

    double ro_j(int j, int M, int N){
        if (j == 1 or j == N-1)
            return 0.25;
        else
            return 1;
    }

double vector_mult(double *w0, double *w1, int M, int N, double hx, double hy){          // Скалярное произведение в пространстве dim(H) =[ 1..(M-1) 1..(N-1) ]
    double res = 0;

    for(int i = 0; i < M+1; i++)
        for (int j = 0; j < N+1; j++)
            if(i == 0 or i == M or j == 0 or j == N)
                res += 0;
            else
                res += ro_i(i, M, N)*ro_j(j, M, N)*w0[i*(N+1) + j]*w1[i*(N+1) + j];

    return res*hx*hy;
}

double norm(double *w, int M, int N, double hx, double hy){                              // Норма вектора без учёта границ
    return sqrt(vector_mult(w, w, M, N, hx, hy));
}

double* B_init(int M, int N, double hx, double hy){                                      // Правая часть уравнения F_ij
    double *res = new double [(M+1)*(N+1)];

    for(int i = 0; i < M+1; i++)
        for (int j = 0; j < N+1; j++){ 
            res[i*(N+1) + j] = func_F(A1 + i*hx, B1 +j*hy);
        }

    return res;
}

double* A_mult(double *w, int M, int N, double hx, double hy){                                // Хитрая функция, результатом которой будет умножение вектора на матрицу А слева
    double *res = new double [(M+1)*(N+1)];
    double wx, wy;

    for(int i = 0; i < M+1; i++)
        for (int j = 0; j < N+1; j++){
            if(i == 0 or i == M or j == 0 or j == N){  
                res[i*(N+1) + j] = w[i*(N+1) + j];         // Граничные условия сохраняем неизменными
            }
            else{
                wx = func_k(A1 + (i+0.5)*hx, B1 + j*hy)*(w[(i+1)*(N+1) + j] - w[i*(N+1) + j])/hx - func_k(A1 + (i-0.5)*hx, B1 + j*hy)*(w[i*(N+1) + j] - w[(i-1)*(N+1) + j])/hx;
                //wx = (w[(i+1)*(N+1) + j] - w[i*(N+1) + j])/hx - (w[i*(N+1) + j] - w[(i-1)*(N+1) + j])/hx;
                wy = func_k(A1 + i*hx, B1 + (j+0.5)*hy)*(w[i*(N+1) + j+1] - w[i*(N+1) + j])/hy - func_k(A1 + i*hx, B1 + (j-0.5)*hy)*(w[i*(N+1) + j] - w[i*(N+1) + j-1])/hy;
                //wy = (w[i*(N+1) + j+1] - w[i*(N+1) + j])/hy - (w[i*(N+1) + j] - w[i*(N+1) + j-1])/hy;
                res[i*(N+1) + j] = -wx/hx - wy/hy + func_q(A1 + i*hx, B1 + j*hy) * w[i*(N+1) + j];
                //res[i*(N+1) + j] = -wx/hx - wy/hy;
            }
    }
                
    return res;
}

int main(int argc, char *argv[])
{	 
    int M, N;
    int i,j;

    ofstream fout("output.txt", ios_base::app);

    if(argc == 3){
         M = atoi(argv[1]);
         N = atoi(argv[2]);
    }
    fout << "M = " << M << " N = " << N << "\n";  

    // Шаги сетки
    double hx = double(A2-A1) / M;
    double hy = double(B2-B1) / N;

    cout << "Start program\n";

    //  Сразу выпрямляем матрицу в вектор и работаем далее с ним
    cout << "Start alloc memory\n";

    double *w = new double [(M+1)*(N+1)];
    //cout << hx << ' ' << hy <<'\n';
    
    cout << "Start initiallizing\n";

    for (i = 0; i < (M+1)*(N+1); i++){            
        w[i] = 0;
    } 

    // Задаём условия на границах 
    // Oстрожно, здесь всё перепутано! 
    for (j = 0; j < M+1; j++){
        w[j] = func_u(A1, B1  + j*hy);                       // Левая
        w[(N+1)*(M+1) - (M+1) + j] = func_u(A2, B1  + j*hy); // Правая
    }
        
    for (i = 0; i < N+1; i++){
        w[i*(M+1)] = func_u(A1 + i*hx, B1);     // Нижняя
        w[i*(M+1) + M] = func_u(A1 + i*hx, B2); // Верхняя
    }
/*
    // Отладочная печать матрицы
    for (i = 0; i < M+1; i++){
        for (j = 0; j < N+1; j++)             
            cout << fixed << w[i*(N+1) + j] << " ";
        cout << endl;
    }
*/
    cout << "Start solving\n";

   // double *r = new double [(M+1)*(N+1)];
   // double *Ar = new double [(M+1)*(N+1)];
    double *w_pr = new double [(M+1)*(N+1)];
    double *B = new double [(M+1)*(N+1)];
    double tau;

    for (i = 0; i < (M+1)*(N+1); i++){            // На первой итерации значение не имеет значения
        w_pr[i] = 1;                                 
    }
    B = B_init(M, N, hx, hy);

    int k = 0;

    // Численный метод решения 
    while (norm(vector_diff(w, w_pr, M, N), M, N, hx, hy) > EPS){
        k++;

        for (i = 0; i < (M+1)*(N+1); i++)           
            w_pr[i] = w[i];

        r = vector_diff(A_mult(w, M, N, hx, hy), B, M, N);

        Ar = A_mult(r, M, N, hx, hy);

        tau = (vector_mult(Ar, r, M, N, hx, hy))/pow(norm(Ar, M, N, hx, hy), 2);
        //cout << fixed << tau << " \n";

        //tau = tau/2;

        for (i = 0; i < (M+1); i++)
            for (j = 0; j < (N+1); j++)
                if(i == 0 or i == M or j == 0 or j == N)
                    continue;
                else
                    w[i*(N+1) + j] = w[i*(N+1) + j] - tau*r[i*(N+1) + j];

        delete[] r;
        delete[] Ar;
         
        //cout << norm(vector_diff(w, w_pr, M, N), M, N, hx, hy) << '\n';
    } 

    for (i = 0; i < M+1; i++){
        for (j = 0; j < N+1; j++)             
            fout << fixed << w[i*(N+1) + j] << " ";
        fout << endl;
    }
    cout << "number of itterations: " << k << '\n';
    cout << "tau: " << tau << '\n';
    cout << "eps: " << EPS << '\n';
     
    cout << "Free memmory\n";

    delete[] w;
    delete[] w_pr;
    delete[] r;
    delete[] Ar;
    delete[] B;
    
    fout.close();

    return 0;
}